import torch
from MyData import MyDataset
from torch.utils.data import DataLoader
from net import Model
from transformers import BertTokenizer
from torch.optim import AdamW  # 使用torch自带的AdamW, BERT 模型的一个优化

#定义训练设备，
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
EPOCH = 100 #训练轮数

token = BertTokenizer.from_pretrained(r"D:\AI\model\bert-base-chinese")

# 直接定义标签映射字典
LABEL_MAP = {
    "开心": 0,
    "伤心": 1,
    "厌恶": 2,
    "惊讶": 3,
    "生气": 4,
    "关心": 5,
    "疑问": 6,
    "平静": 7
}
# 自定义函数，对数据进行编码处理
def collate_fn(data):
    sentes = [i[0] for i in data] # 获取所有句子
    labels = [i[1] for i in data] # id
    #编码
    data = token.batch_encode_plus(batch_text_or_text_pairs=sentes,
                                   truncation=True,
                                   padding="max_length",
                                   max_length=128,  #按照每条文本最长的，但是太大的话，会报错。需要改模型的编码长度 position_embeddings。例子:/test02.py
                                   return_tensors="pt",
                                   return_length=True
                                   )
    input_ids = data["input_ids"]
    attention_mask = data["attention_mask"]
    token_type_ids = data["token_type_ids"]

    # 将标签转换为数值并创建张量
    # 如果标签是字符串类别，需要先映射为数字
    if isinstance(labels[0], str):
        # 创建标签到数字的映射
        # unique_labels = list(set(labels))
        # label_to_id = {label: idx for idx, label in enumerate(unique_labels)}
        # labels = [label_to_id[label] for label in labels]
        labels = [LABEL_MAP[label] for label in labels]
        # print(f"标签映射: {labels}")  # 调试信息

    labels = torch.LongTensor(labels) # 转换成张量
    return input_ids, attention_mask, token_type_ids, labels

#创建数据集
train_dataset = MyDataset("train")
#创建数据加载器
train_loader = DataLoader(dataset=train_dataset,
                          batch_size=32, #每次加载32个数据
                          shuffle=True, #把数据集打乱
                          drop_last=True, #
                          collate_fn=collate_fn #自定义方法，对数据进行编码处理
                          ) #


if __name__ == '__main__':
    # 开始模型训练
    print(DEVICE)

    # 先获取一批数据来确定类别数
    sample_data = next(iter(train_loader))
    num_classes = len(torch.unique(sample_data[3]))
    print(f"检测到类别数: {num_classes}")

    # 确保模型输出层与类别数匹配
    model = Model().to(DEVICE) #模型
    optimizer = AdamW(model.parameters(), lr=5e-4) # 使用torch自带的优化器AdamW, BERT 模型的一个优化
    loss_func = torch.nn.CrossEntropyLoss() # 损失函数

    model.train()
    for epoch in range(EPOCH): #训练多轮
        for step, (input_ids, attention_mask, token_type_ids, labels) in enumerate(train_loader): #从数据加载器中获取数据
            #将数据放到设备上
            input_ids, attention_mask, token_type_ids, labels = input_ids.to(DEVICE), attention_mask.to(DEVICE), token_type_ids.to(DEVICE), labels.to(DEVICE)
            #执行前向计算得到输出
            output = model(input_ids, attention_mask, token_type_ids)

            # 调试信息
            if step == 0 and epoch == 0:
                print(f"模型输出形状: {output.shape}")
                print(f"标签形状: {labels.shape}")
                print(f"标签唯一值: {torch.unique(labels)}")

            loss = loss_func(output, labels)

            optimizer.zero_grad() # 清空梯度 , 如果是验证，就不需要了
            loss.backward() # 反向传播, 如果是验证，就不需要了
            optimizer.step() #  更新参数, 如果是验证，就不需要了

            if step % 5 == 0: #每5个step打印一次结果
                out = output.argmax(dim=1)
                acct = (out == labels).sum().item()/len(labels)
                print(f"epoch:{epoch}, step:{step}, loss:{loss.item()}, acc:{acct}")

        #保存模型参数
        torch.save(model.state_dict(), f"params/{epoch}bert.pth")
        print(epoch, "保存成功")


